import numpy as np
import pygame

# 初始化Pygame
pygame.init()

# 屏幕尺寸
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Gravitational Simulation")

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

# 时间步长
dt = 0.01

# 万有引力常数
G = 6.67430e-11

class Vector:
    """一个简单的二维向量类"""
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __add__(self, other):
        return Vector(self.x + other.x, self.y + other.y)

    def __sub__(self, other):
        return Vector(self.x - other.x, self.y - other.y)

    def __mul__(self, scalar):
        return Vector(self.x * scalar, self.y * scalar)

    def __truediv__(self, scalar):
        return Vector(self.x / scalar, self.y / scalar)

    def magnitude(self):
        return np.sqrt(self.x ** 2 + self.y ** 2)

    def normalize(self):
        mag = self.magnitude()
        return Vector(self.x / mag, self.y / mag)

class Planet:
    """行星类，包含位置、速度、加速度和质量"""
    def __init__(self, pos, vel, mass):
        self.pos = pos
        self.vel = vel
        self.acc = Vector(0, 0)
        self.mass = mass

def calculate_gravitational_force(planet1, planet2):
    """计算两个天体之间的万有引力"""
    r = planet2.pos - planet1.pos
    distance = r.magnitude()
    force_magnitude = G * planet1.mass * planet2.mass / distance ** 2
    force_direction = r.normalize()
    return force_direction * force_magnitude

def runge_kutta_4(planet, dt, forces):
    """使用四阶龙格库塔法更新行星状态"""
    k1_vel = forces(planet) / planet.mass * dt
    k1_pos = planet.vel * dt

    k2_vel = forces(Planet(planet.pos + 0.5 * k1_pos, planet.vel + 0.5 * k1_vel, planet.mass)) / planet.mass * dt
    k2_pos = (planet.vel + 0.5 * k1_vel) * dt

    k3_vel = forces(Planet(planet.pos + 0.5 * k2_pos, planet.vel + 0.5 * k2_vel, planet.mass)) / planet.mass * dt
    k3_pos = (planet.vel + 0.5 * k2_vel) * dt

    k4_vel = forces(Planet(planet.pos + k3_pos, planet.vel + k3_vel, planet.mass)) / planet.mass * dt
    k4_pos = (planet.vel + k3_vel) * dt

    planet.pos = planet.pos + (k1_pos + 2 * k2_pos + 2 * k3_pos + k4_pos) / 6
    planet.vel = planet.vel + (k1_vel + 2 * k2_vel + 2 * k3_vel + k4_vel) / 6

def main():
    # 创建行星
    planets = [
        Planet(Vector(200, 300), Vector(0, 0.5), 5.972e24),  # 地球
        Planet(Vector(500, 300), Vector(0, -0.5), 7.348e22)  # 月球
    ]

    clock = pygame.time.Clock()
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        def total_force(planet):
            force = Vector(0, 0)
            for other in planets:
                if other != planet:
                    force += calculate_gravitational_force(planet, other)
            return force

        for planet in planets:
            runge_kutta_4(planet, dt, total_force)

        screen.fill(BLACK)
        for planet in planets:
            pygame.draw.circle(screen, WHITE, (int(planet.pos.x), int(planet.pos.y)), 10)

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()